from .bases import BaseImageDataset
import os


class NAIC(BaseImageDataset):

    def __init__(self):
        super().__init__()
        self.root = '/Users/zmh/DataFiles/Projects/PyCharm/reid/res/train'
        self.train_feature = os.path.join(self.root, 'train_feature')
        f = open(os.path.join(self.root, 'train_list.txt'), 'r')
        train_data = []
        for l in f.readlines():
            dat_path, person_id = l.split(' ')
            person_id = int(person_id)
            train_data.append((os.path.join(self.train_feature, dat_path), person_id, 0))
        self.train = train_data[0:-8]
        self.query = train_data[-8:]
        self.gallery = train_data
        self.num_train_pids, self.num_train_imgs, self.num_train_cams = self.get_imagedata_info(self.train)
        self.num_query_pids, self.num_query_imgs, self.num_query_cams = self.get_imagedata_info(self.query)
        self.num_gallery_pids, self.num_gallery_imgs, self.num_gallery_cams = self.get_imagedata_info(self.gallery)
